package edu.zju.gis.dbfg.model.statistic;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.datasource.DataReaderRDD;
import edu.zju.gis.dbfg.model.util.CrsPick;
import edu.zju.gis.dbfg.model.util.GeoItem;
import edu.zju.gis.dbfg.model.util.TileID;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.api.java.function.PairFunction;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.Geometry;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.io.File;
import java.io.FileWriter;

import static edu.zju.gis.dbfg.model.base.Constant.DEFAULT_CRS;
import static edu.zju.gis.dbfg.model.base.Constant.DEFAULT_PYRAMID_CONFIG;

public class AbnormalDetect {

    private final static Logger logger = LoggerFactory.getLogger(AbnormalDetect.class);

    public static void main (String[] args) throws Exception{

        //Setup environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
        conf.setMaster("local[4]");
        conf.setAppName("AbnormalDetect");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        //Setup model arguments
        logger.info("Setup arguments");
        AbnormalDetectArgs mArgs = AbnormalDetectArgs.initArgs(args);
        if(mArgs == null){
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(AbnormalDetect.class, "main()", msg, args);
        }

        CoordinateReferenceSystem crs1 = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs1()));
        CoordinateReferenceSystem crs2 = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs2()));
        CoordinateReferenceSystem targetCRS = DEFAULT_CRS();
        PyramidConfig targetPyramidConfig = DEFAULT_PYRAMID_CONFIG();

        MathTransform transform1 = null;
        MathTransform transform2 = null;
        if(!crs1.equals(targetCRS)){
            transform1 = CRS.findMathTransform(crs1, targetCRS);
        }
        if(!crs2.equals(targetCRS)){
            transform2 = CRS.findMathTransform(crs2, targetCRS);
        }

        //Read data1 from datasource and map it to geoItem
        String input1 = mArgs.getInput1();
        String type1 = mArgs.getType1();
        String pac1 = mArgs.getPacIndex1();
        String area1 = mArgs.getAreaIndex1();
        JavaRDD<GeoItem> gi1 = null;
        Map<String, Object> attrs1 = new HashMap<>();
        attrs1.put("PAC", pac1);
        //如指定了面积字段，则不需要重新计算，直接读取
        if(!area1.equals("-1")) //Modify...
        {
            attrs1.put("AREA", area1);
            gi1 = DataReaderRDD.readToRDD(type1, jsc, input1, null, attrs1, transform1);
        }
        //如未指定面积字段，则重新计算面积并记录
        else{
            JavaRDD<GeoItem> origi = DataReaderRDD.readToRDD(type1, jsc, input1, null, attrs1, transform1);
            gi1 = origi.map(new Function<GeoItem, GeoItem>() {
                @Override
                public GeoItem call(GeoItem geoItem) throws Exception {
                    Geometry geometry = geoItem.getGeometry();
                    Double area = geometry.getArea();
                    GeoItem newgi = new GeoItem();
                    newgi.setFid(geoItem.getFid());
                    newgi.setGeometry(geometry);
                    Map<String,Object> props = new HashMap<>();
                    props.put("AREA", area);
                    newgi.setProp(props);
                    return newgi;
                }
            });
        }
        //将原始数据映射为以PAC为key
        JavaPairRDD<String, GeoItem> gi1Pac = gi1.mapToPair(new PairFunction<GeoItem, String, GeoItem>() {
            @Override
            public Tuple2<String, GeoItem> call(GeoItem input) throws Exception {
                String pac = (String)input.getProp().get("PAC");
                return new Tuple2<>(pac, input);
            }
        });
        gi1Pac.cache();
        //根据PAC汇总AREA
        JavaPairRDD<String,GeoItem> gi1PairAreaSum = gi1Pac.reduceByKey(new Function2<GeoItem, GeoItem, GeoItem>() {
            @Override
            public GeoItem call(GeoItem gi1, GeoItem gi2) throws Exception {
                GeoItem result = new GeoItem();
                Double area = (Double) gi1.getProp().get("AREA") + (Double) gi2.getProp().get("AREA");
                result.setFid(gi1.getFid());
                result.setProp(gi1.getProp());
                result.getProp().put("AREA", area);
                return result;
            }
        });
        //根据PAC，将汇总后的AREA添加到原始图斑中
        JavaPairRDD<String, Tuple2<GeoItem, GeoItem>> gi1PairPacArea = gi1Pac.join(gi1PairAreaSum);
        JavaRDD<GeoItem> gi1PacArea = gi1PairPacArea.map(new Function<Tuple2<String, Tuple2<GeoItem, GeoItem>>, GeoItem>() {
            @Override
            public GeoItem call(Tuple2<String, Tuple2<GeoItem, GeoItem>> input) throws Exception {
                GeoItem result = new GeoItem();
                GeoItem gi1 = input._2._1();  //原始图斑
                GeoItem gi2 = input._2._2();  //按PAC汇总面积后的图斑
                result.setFid(gi1.getFid());
                Map<String, Object> props = gi1.getProp();
                props.put("TOTAL_AREA",gi2.getProp().get("AREA"));
                result.setProp(props);
                result.setGeometry(gi1.getGeometry());
                return result;
            }
        });
        //按格网号组织
        JavaPairRDD<String, GeoItem> gi1Pair = DataReaderRDD.mapToPair(gi1PacArea, targetPyramidConfig);

        // Read data2 from datasource and map it to geoItem
        String input2 = mArgs.getInput2();
        String type2 = mArgs.getType2();
        String cc2 = mArgs.getCcIndex2();
        String pac2 = mArgs.getPacIndex2();
        String area2 = mArgs.getAreaIndex2();
        Map<String, Object> attrs2 = new HashMap<>();
        attrs2.put("CC", cc2);
        attrs2.put("PAC", pac2);
        attrs2.put("AREA", area2);
        JavaPairRDD<String, GeoItem> gi2Pair = null;

        if(type2.equals("ES")){
            gi1Pair.cache();
            List<TileID> tileIDs = gi1Pair.map(t-> TileID.fromString(t._1)).distinct().collect();
            gi2Pair = DataReaderRDD.readToPairRDD(type2, jsc, input2, attrs2, tileIDs, targetPyramidConfig);
        }else{
            JavaRDD<GeoItem> gi2 = DataReaderRDD.readToRDD(type2, jsc, input2, null, attrs2, transform2);
            gi2Pair = DataReaderRDD.mapToPair(gi2, targetPyramidConfig);
        }

        //Intersect
        JavaPairRDD<String, Tuple2<GeoItem, GeoItem>> joined = gi1Pair.join(gi2Pair);
        JavaPairRDD<String, GeoItem> tempRDD = joined.mapToPair(new PairFunction<Tuple2<String, Tuple2<GeoItem, GeoItem>>, String, GeoItem>() {
            @Override
            public Tuple2<String, GeoItem> call(Tuple2<String, Tuple2<GeoItem, GeoItem>> input) throws Exception {
                GeoItem gi1 = input._2._1;
                GeoItem gi2 = input._2._2;
                Geometry g1 = gi1.getGeometry(); //Extent图斑
                Geometry g2 = gi2.getGeometry(); //LCRA图斑
                if(!g1.intersects(g2)){
                    return new Tuple2<>("EMPTY",null);
                }
                Geometry intersection = g1.intersection(g2);
                Double area = intersection.getArea();
                String cc = (String)gi2.getProp().get("CC");
                String pac = (String)gi2.getProp().get("PAC");
                String id = pac+"##"+cc;
                GeoItem result = new GeoItem();
                result.setFid(id);
                Map<String, Object> props = new HashMap<>();
                props.put("CC",cc);
                props.put("PAC",pac);
                props.put("AREA", area);  //相交几何面积
                props.put("TOTAL_AREA",gi1.getProp().get("TOTAL_AREA"));  //Extent要素按县汇总的原始面积
                result.setProp(props);
                result.setGeometry(intersection);
                return new Tuple2<>(id,result);
            }
        }).filter(new Function<Tuple2<String, GeoItem>, Boolean>(){
            @Override
                    public Boolean call(Tuple2<String, GeoItem> input)  throws Exception{
                return !input._1.equals("EMPTY");
            }
        });

        tempRDD.cache();
        writeToFile(tempRDD, "D:\\ZJU\\test\\FX\\1206\\abdIntersect.wkt");

        //Sum up area
        JavaPairRDD<String, GeoItem> resultRDD = tempRDD.reduceByKey(new Function2<GeoItem, GeoItem, GeoItem>() {
            @Override
            public GeoItem call(GeoItem gi1, GeoItem gi2) throws Exception {
                Double area = (Double) gi1.getProp().get("AREA")+(Double) gi2.getProp().get("AREA");
                GeoItem result = new GeoItem();
                result.setFid(gi1.getFid());
                result.setProp(gi1.getProp());
                result.getProp().put("AREA", area);
                return result;
            }
        });


        //Write out
        Map<String, GeoItem> result = resultRDD.collectAsMap();
        File file = new File(mArgs.getOutput());
        if (file.exists()) {
            file.delete();
        }
        FileWriter writer = new FileWriter(file);
        writer.write(String.format("%s\t%s\t%s\t%s\t%s\n",pac2, cc2, area2, "totalArea", "ratio"));
        for (String key: result.keySet()) {
            String[] fields = key.split("##");
            String cc = fields[1];
            String pac = fields[0];
            GeoItem gi = result.get(key);
            Double area = (Double) gi.getProp().get("AREA");
            Double totalArea = (Double) gi.getProp().get("TOTAL_AREA");
            Double ratio = area/totalArea;

            String s = String.format("%s\t%s\t%.9f\t%.9f\t%.9f", pac, cc, area, totalArea, ratio);
            writer.write(s+"\n");
        }
        writer.flush();
        writer.close();
    }

    public static void writeToFile(JavaPairRDD<String, GeoItem> rdd, String filename) throws IOException {
        List<Tuple2<String, GeoItem>> featureResult = rdd.collect();
        File featureFile = new File(filename);
        if (featureFile.exists()) featureFile.delete();
        FileWriter featureWriter = new FileWriter(featureFile);
        for(Tuple2<String, GeoItem> feature: featureResult) {
            GeoItem gi = feature._2;
            String id = gi.getFid();
            Geometry geometry = gi.getGeometry();
            String wkt = geometry.toText();
            featureWriter.write(String.format("%s\t%s\n", id, wkt));
        }
        featureWriter.flush();
        featureWriter.close();

}}
